// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include <unordered_map>
#include <unordered_set>
#include <vector>

#include <tbir/ir/hlo_builtin.h>
#include <tbir/ir/stmt_functor.h>

namespace tbir::ir {

    class FuseContCasterOptimizer : public StmtExprMutator {
    public:
        BaseFunc run(const BaseFunc &func) {
            return this->VisitExpr(func);
        }

        HLOExpr VisitExpr_(const HLOCastNode *op) override {
            const auto &op_type = RemoveReference(op->checked_type());
            auto raw_value = RemoveNestedCaster(op->value);
            if (raw_value.same_as(op->value)) {
                return runtime::GetRef<HLOExpr>(op);
            }
            return HLOCast(op->checked_type(), raw_value, op->span);
        }

        PrimExpr VisitExpr_(const HLOCastPrimNode *op) override {
            const auto &op_type = RemoveReference(op->checked_type());
            auto raw_value = RemoveNestedCaster(op->value);
            if (raw_value.same_as(op->value)) {
                return runtime::GetRef<PrimExpr>(op);
            }
            return HLOCastPrim(op->dtype, raw_value, op->span);
        }

    protected:
        static BaseExpr RemoveNestedCaster(BaseExpr e) {
            if (auto *cast_node = e.as<HLOCastNode>()) {
                return RemoveNestedCaster(cast_node->value);
            }
            if (auto *cast_node = e.as<HLOCastPrimNode>()) {
                return RemoveNestedCaster(cast_node->value);
            }
            return e;
        }
    };

}  // namespace tbir::ir
